#!/usr/bin/env python3
"""
自动更新 requirements.txt 中的包版本到最新版本

用途：
- 解析 requirements.txt 文件中的包名
- 从 PyPI 获取每个包的最新版本
- 更新版本号并写入文件
- 处理特殊版本约束（>=, ==, ~=）

使用方法：
python scripts/update_requirements.py [--dry-run] [--backup]

参数：
- --dry-run: 只显示将要更新的版本，不实际写入文件
- --backup: 更新前备份原文件
"""

import argparse
import json
import re
import shutil
import sys
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import requests


class RequirementsUpdater:
    """requirements.txt 版本更新器"""
    
    # 支持的 PyPI 源配置
    PYPI_SOURCES = {
        'official': {
            'name': 'PyPI 官方源',
            'pypi_url': 'https://pypi.org/pypi',
            'pip_index': 'https://pypi.org/simple'
        },
        'tsinghua': {
            'name': '清华大学源',
            'pypi_url': 'https://pypi.tuna.tsinghua.edu.cn/pypi',
            'pip_index': 'https://pypi.tuna.tsinghua.edu.cn/simple'
        },
        'aliyun': {
            'name': '阿里云源',
            'pypi_url': 'https://mirrors.aliyun.com/pypi',
            'pip_index': 'https://mirrors.aliyun.com/pypi/simple'
        },
        'douban': {
            'name': '豆瓣源',
            'pypi_url': 'https://pypi.douban.com/pypi',
            'pip_index': 'https://pypi.douban.com/simple'
        },
        'ustc': {
            'name': '中科大源',
            'pypi_url': 'https://pypi.mirrors.ustc.edu.cn/pypi',
            'pip_index': 'https://pypi.mirrors.ustc.edu.cn/simple'
        }
    }
    
    def __init__(self, requirements_file: str = "requirements.txt", source: str = "official"):
        """
        初始化更新器
        
        参数:
            requirements_file: requirements.txt 文件路径
            source: PyPI 源名称，可选值: official, tsinghua, aliyun, douban, ustc
        """
        self.requirements_file = Path(requirements_file)
        
        # 设置 PyPI 源
        if source not in self.PYPI_SOURCES:
            print(f"⚠️  未知的源: {source}，使用默认的官方源")
            source = "official"
        
        self.source_config = self.PYPI_SOURCES[source]
        self.pypi_url = self.source_config['pypi_url']
        self.pip_index = self.source_config['pip_index']
        
        print(f"🌐 使用 PyPI 源: {self.source_config['name']}")
        
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'requirements-updater/1.0'
        })
    
    def parse_requirement_line(self, line: str) -> Optional[Tuple[str, str, str]]:
        """
        解析 requirements 行
        
        参数:
            line: requirements.txt 中的一行
            
        返回:
            (包名, 操作符, 版本号) 或 None
        """
        line = line.strip()
        
        # 跳过注释和空行
        if not line or line.startswith('#'):
            return None
        
        # 匹配包名和版本约束
        # 支持格式: package>=1.0.0, package==1.0.0, package~=1.0.0, package[extra]>=1.0.0
        pattern = r'^([a-zA-Z0-9\-_\[\]]+)\s*([><=~!]+)\s*([0-9\.]+[a-zA-Z0-9\-_\.]*)'
        match = re.match(pattern, line)
        
        if match:
            package_name = match.group(1)
            operator = match.group(2)
            version = match.group(3)
            
            # 处理包名中的 extras（如 uvicorn[standard]）
            if '[' in package_name:
                package_name = package_name.split('[')[0]
            
            return package_name, operator, version
        
        return None
    
    def get_latest_version(self, package_name: str) -> Optional[str]:
        """
        从 PyPI 获取包的最新版本
        
        参数:
            package_name: 包名
            
        返回:
            最新版本号或 None
        """
        try:
            url = f"{self.pypi_url}/{package_name}/json"
            response = self.session.get(url, timeout=10)
            response.raise_for_status()
            
            data = response.json()
            latest_version = data['info']['version']
            
            print(f"  {package_name}: 最新版本 {latest_version}")
            return latest_version
            
        except requests.exceptions.RequestException as e:
            print(f"  ❌ 获取 {package_name} 版本失败: {e}")
            return None
        except (KeyError, json.JSONDecodeError) as e:
            print(f"  ❌ 解析 {package_name} 版本数据失败: {e}")
            return None
    
    def update_requirements(self, dry_run: bool = False, backup: bool = True) -> bool:
        """
        更新 requirements.txt 文件
        
        参数:
            dry_run: 是否只显示变更而不实际写入
            backup: 是否备份原文件
            
        返回:
            是否成功更新
        """
        if not self.requirements_file.exists():
            print(f"❌ 文件不存在: {self.requirements_file}")
            return False
        
        print(f"📖 读取文件: {self.requirements_file}")
        
        # 读取原文件
        with open(self.requirements_file, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        
        updated_lines = []
        updates_made = []
        
        print("\n🔍 检查包版本...")
        
        for line in lines:
            original_line = line
            parsed = self.parse_requirement_line(line)
            
            if parsed:
                package_name, operator, current_version = parsed
                print(f"\n📦 检查 {package_name} (当前: {current_version})")
                
                latest_version = self.get_latest_version(package_name)
                
                if latest_version and latest_version != current_version:
                    # 构建新的行
                    # 保持原有的操作符，除非是 == 则改为 >=
                    new_operator = ">=" if operator == "==" else operator
                    
                    # 重建包名（包含可能的 extras）
                    original_package = line.split(operator)[0].strip()
                    new_line = f"{original_package}{new_operator}{latest_version}\n"
                    
                    updated_lines.append(new_line)
                    updates_made.append({
                        'package': package_name,
                        'from': current_version,
                        'to': latest_version,
                        'operator': new_operator
                    })
                    
                    print(f"  ✅ 将更新: {current_version} -> {latest_version}")
                else:
                    updated_lines.append(original_line)
                    if latest_version:
                        print(f"  ✓ 已是最新版本")
            else:
                # 保持注释和空行不变
                updated_lines.append(original_line)
        
        # 显示更新摘要
        if updates_made:
            print(f"\n📋 发现 {len(updates_made)} 个包需要更新:")
            for update in updates_made:
                print(f"  • {update['package']}: {update['from']} -> {update['to']}")
        else:
            print("\n✅ 所有包都是最新版本！")
            return True
        
        if dry_run:
            print("\n🔍 预览模式 - 不会实际修改文件")
            print("\n更新后的内容预览:")
            print("=" * 50)
            for line in updated_lines[:20]:  # 只显示前20行
                print(line.rstrip())
            if len(updated_lines) > 20:
                print("... (还有更多行)")
            return True
        
        # 备份原文件
        if backup:
            backup_file = self.requirements_file.with_suffix(
                f'.txt.backup_{datetime.now().strftime("%Y%m%d_%H%M%S")}'
            )
            shutil.copy2(self.requirements_file, backup_file)
            print(f"\n💾 已备份原文件到: {backup_file}")
        
        # 写入更新后的内容
        try:
            with open(self.requirements_file, 'w', encoding='utf-8') as f:
                f.writelines(updated_lines)
            
            print(f"\n✅ 成功更新 {self.requirements_file}")
            print(f"   更新了 {len(updates_made)} 个包的版本")
            print(f"   使用的源: {self.source_config['name']}")
            
            return True
            
        except Exception as e:
            print(f"\n❌ 写入文件失败: {e}")
            return False


def show_sources():
    """显示所有可用的 PyPI 源"""
    print("🌐 可用的 PyPI 源:")
    print("=" * 40)
    for key, config in RequirementsUpdater.PYPI_SOURCES.items():
        print(f"  {key:10} - {config['name']}")
        print(f"             {config['pypi_url']}")
        print()


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="自动更新 requirements.txt 中的包版本到最新版本",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  python scripts/update_requirements.py                    # 直接更新（使用官方源）
  python scripts/update_requirements.py --dry-run         # 预览更新
  python scripts/update_requirements.py --source official # 使用官方源
  python scripts/update_requirements.py --source aliyun   # 使用阿里云源
  python scripts/update_requirements.py --backup          # 更新前备份
        """
    )
    
    parser.add_argument(
        '--dry-run', 
        action='store_true',
        help='只显示将要更新的版本，不实际写入文件'
    )
    
    parser.add_argument(
        '--backup',
        action='store_true', 
        default=True,
        help='更新前备份原文件（默认启用）'
    )
    
    parser.add_argument(
        '--no-backup',
        action='store_true',
        help='不备份原文件'
    )
    
    parser.add_argument(
        '--file',
        default='requirements.txt',
        help='指定 requirements 文件路径（默认: requirements.txt）'
    )
    
    parser.add_argument(
        '--source',
        choices=['official', 'tsinghua', 'aliyun', 'douban', 'ustc'],
        default='official',
        help='指定 PyPI 源 (默认: official)'
    )
    
    parser.add_argument(
        '--list-sources',
        action='store_true',
        help='显示所有可用的 PyPI 源'
    )
    
    args = parser.parse_args()
    
    # 处理备份选项
    backup = args.backup and not args.no_backup
    
    # 如果只是显示源列表，直接显示并退出
    if args.list_sources:
        show_sources()
        return
    
    print("🚀 Requirements 版本更新工具")
    print("=" * 40)
    
    try:
        updater = RequirementsUpdater(args.file, source=args.source)
        success = updater.update_requirements(
            dry_run=args.dry_run,
            backup=backup
        )
        
        if success:
            print("\n🎉 操作完成！")
            if not args.dry_run:
                print("\n💡 提示: 更新后建议运行以下命令重新安装依赖:")
                print(f"   pip install -i {updater.pip_index} -r requirements.txt --upgrade")
        else:
            print("\n❌ 操作失败！")
            sys.exit(1)
            
    except KeyboardInterrupt:
        print("\n\n⏹️  操作被用户取消")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 发生意外错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main() 